#include "RHIRenderPass.h"
#include "../../Utils/LogUtil.h"
#include "../../Core/CoreGlobals.h"
#include "../RenderState.h"
#include <algorithm>
#include "../../Game/Process/RenderProcess.h"

RHIRenderPass::BuilderDetails::BuilderDetails()
{
	srcStageMask = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
	srcAccessMask = VK_ACCESS_MEMORY_READ_BIT;

	dstStageMask = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
	dstAccessMask = VK_ACCESS_MEMORY_READ_BIT;
}

RHIRenderPass::Builder::Builder() noexcept = default;
RHIRenderPass::Builder::~Builder() noexcept = default;
RHIRenderPass::Builder::Builder(RHIRenderPass::Builder const& rhs) noexcept = default;
RHIRenderPass::Builder::Builder(RHIRenderPass::Builder && rhs) noexcept = default;
RHIRenderPass::Builder& RHIRenderPass::Builder::operator=(RHIRenderPass::Builder const& rhs) noexcept = default;
RHIRenderPass::Builder& RHIRenderPass::Builder::operator=(RHIRenderPass::Builder && rhs) noexcept = default;


RHIRenderPass::Builder& RHIRenderPass::Builder::AddRenderTarget(RefCountPtr<RenderTarget> renderTarget) noexcept
{
	mImpl->renderTargets.push_back(renderTarget);
	return *this;
}

RHIRenderPass::Builder& RHIRenderPass::Builder::AddPipeline(RefCountPtr<GraphicsPipeline> pipeline) noexcept
{
	mImpl->pipelines.push_back(pipeline);
	return *this;
}

RHIRenderPass::Builder& RHIRenderPass::Builder::SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept
{
	mImpl->srcStageMask = srcStageMask;
	return *this;
}

RHIRenderPass::Builder& RHIRenderPass::Builder::SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept
{
	mImpl->srcAccessMask = srcAccessMask;
	return *this;
}

RHIRenderPass::Builder& RHIRenderPass::Builder::SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept
{
	mImpl->dstStageMask = dstStageMask;
	return *this;
}

RHIRenderPass::Builder& RHIRenderPass::Builder::SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept
{
	mImpl->dstAccessMask = dstAccessMask;
	return *this;
}

RefCountPtr<RHIRenderPass> RHIRenderPass::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<RHIRenderPass>(NewObject<RHIRenderPass>(*mImpl));
}

RefCountPtr<RHIRenderPass> RHIRenderPass::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	RHIRenderPass* renderPass = NewObject<RHIRenderPass>(*mImpl);
	process->AddRenderPass(name, renderPass);
	return RefCountPtr<RHIRenderPass>(renderPass);
}

void RHIRenderPass::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveRenderPass(mName, this);
		}
		DeleteObject<RHIRenderPass>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveRenderPass(mName, this);
			}
			DeleteObject<RHIRenderPass>(this);
			});
	}
}

RHIRenderPass::~RHIRenderPass()
{
	GRenderState->DelayGCRenderPass(mRenderPass);
}

VkRenderPass RHIRenderPass::GetRenderPass() const
{
	return mRenderPass;
}

const std::vector<std::vector<VkImageView>>& RHIRenderPass::GetRenderTargetViewMap() const
{
	return mRenderTargetViewMap;
}

const std::vector<VkClearValue>& RHIRenderPass::GetClearValues() const
{
	return mClearValues;
}

uint32 RHIRenderPass::GetWidth() const
{
	return mWidth;
}

uint32 RHIRenderPass::GetHeight() const
{
	return mHeight;
}

uint32 RHIRenderPass::GetRenderTargetSize() const
{
	return mRenderTargetSize;
}

uint32 RHIRenderPass::GetRenderTargetIndex(RefCountPtr<RenderTarget> renderTarget) const
{
	for (uint32 i = 0; i < mRenderTargets.size(); ++i)
	{
		if (mRenderTargets[i] == renderTarget)
		{
			return i;
		}
	}
	// no renderTarget exist
	// make sure the renderTarget is add before make subpass
	assert(0);
	return 0;
}

RHIRenderPass::RHIRenderPass(BuilderDetails & builderDetails)
{
	mRenderTargets.assign(builderDetails.renderTargets.begin(), builderDetails.renderTargets.end());
	mRenderTargetSize = static_cast<uint32>(mRenderTargets.size());

	mPipelines.assign(builderDetails.pipelines.begin(), builderDetails.pipelines.end());

	mSrcStageMask = builderDetails.srcStageMask;
	mSrcAccessMask = builderDetails.srcAccessMask;

	mDstStageMask = builderDetails.dstStageMask;
	mDstAccessMask = builderDetails.dstAccessMask;



	// Check data vaild, all renderTarget's width and height should be same

	mClearValues.resize(mRenderTargets.size(), {});

	mRenderTargetViewMap.resize(GRenderState->GetFrameCount(), {});

	std::vector<VkAttachmentDescription> attachmentDescriptions;
	attachmentDescriptions.resize(mRenderTargets.size(), {});

	for (uint32 i = 0; i < mRenderTargets.size(); ++i)
	{
		if (i == 0)
		{
			mWidth = mRenderTargets[i]->GetWitdh();
			mHeight = mRenderTargets[i]->GetHeight();
		}
		else
		{
			assert(mWidth == mRenderTargets[i]->GetWitdh());
			assert(mHeight == mRenderTargets[i]->GetHeight());
		}
		mClearValues[i] = mRenderTargets[i]->GetClearValue();
		attachmentDescriptions[i] = mRenderTargets[i]->GetAttachmentDescription();

		// Get renderTargetView
		for (uint32 j = 0; j < GRenderState->GetFrameCount(); ++j)
		{
			mRenderTargetViewMap[j].push_back(mRenderTargets[i]->GetImageView(j));
		}
	}



	std::vector<SubpassData> subpassDatas;
	std::vector<VkSubpassDescription> subpassDescriptions;
	subpassDatas.resize(mPipelines.size(), SubpassData());
	subpassDescriptions.resize(mPipelines.size(), {});
	for (uint32 i = 0; i < mPipelines.size(); ++i)
	{
		subpassDescriptions[i].pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;

		subpassDatas[i].colorAttachments.resize(mPipelines[i]->mColorRenderTargets.size(), {});
		for (uint32 j = 0; j < mPipelines[i]->mColorRenderTargets.size(); ++j)
		{
			subpassDatas[i].colorAttachments[j] = { GetRenderTargetIndex(mPipelines[i]->mColorRenderTargets[j]), mPipelines[i]->mColorImageLayouts[j] };
		}
		std::sort(std::begin(subpassDatas[i].colorAttachments), std::end(subpassDatas[i].colorAttachments),
			[](const VkAttachmentReference& a, const VkAttachmentReference& b) {
				return a.attachment < b.attachment; });
		subpassDescriptions[i].colorAttachmentCount = static_cast<uint32>(mPipelines[i]->mColorRenderTargets.size());
		subpassDescriptions[i].pColorAttachments = subpassDatas[i].colorAttachments.data();




		if (mPipelines[i]->mDepthRenderTarget.IsValid())
		{
			subpassDatas[i].depthStencilAttachment = { GetRenderTargetIndex(mPipelines[i]->mDepthRenderTarget), mPipelines[i]->mDepthImageLayout };
			subpassDescriptions[i].pDepthStencilAttachment = &subpassDatas[i].depthStencilAttachment;
		}




		if (mPipelines[i]->mResolveRenderTarget.IsValid())
		{
			subpassDatas[i].resolveAttachment = { GetRenderTargetIndex(mPipelines[i]->mResolveRenderTarget), mPipelines[i]->mResolveImageLayout };
			subpassDescriptions[i].pResolveAttachments = &subpassDatas[i].resolveAttachment;
		}




		subpassDatas[i].inputAttachments.resize(mPipelines[i]->mInputRenderTargets.size(), {});
		for (uint32 j = 0; j < mPipelines[i]->mInputRenderTargets.size(); ++j)
		{
			subpassDatas[i].inputAttachments[j] = { GetRenderTargetIndex(mPipelines[i]->mInputRenderTargets[j]), mPipelines[i]->mInputImageLayouts[j] };
		}
		std::sort(std::begin(subpassDatas[i].inputAttachments), std::end(subpassDatas[i].inputAttachments),
			[](const VkAttachmentReference& a, const VkAttachmentReference& b) {
				return a.attachment < b.attachment; });
		subpassDescriptions[i].inputAttachmentCount = static_cast<uint32>(mPipelines[i]->mInputRenderTargets.size());
		subpassDescriptions[i].pInputAttachments = subpassDatas[i].inputAttachments.data();
	}


	std::vector<VkSubpassDependency> subpassDependencys;
	subpassDependencys.resize(mPipelines.size() + 1, { 0, 0, 0, 0, 0, 0, VK_DEPENDENCY_BY_REGION_BIT });

	subpassDependencys[0].srcSubpass = VK_SUBPASS_EXTERNAL;
	subpassDependencys[0].dstSubpass = 0;
	subpassDependencys[0].srcStageMask = mSrcStageMask;
	subpassDependencys[0].srcAccessMask = mSrcAccessMask;

	subpassDependencys[mPipelines.size()].srcSubpass = 0;
	subpassDependencys[mPipelines.size()].dstSubpass = VK_SUBPASS_EXTERNAL;
	subpassDependencys[mPipelines.size()].dstStageMask = mDstStageMask;
	subpassDependencys[mPipelines.size()].dstAccessMask = mDstAccessMask;

	for (uint32 i = 0; i < mPipelines.size(); ++i)
	{
		subpassDependencys[i].dstStageMask = mPipelines[i]->mDstStageMask;
		subpassDependencys[i].dstAccessMask = mPipelines[i]->mDstAccessMask;

		subpassDependencys[i + 1].srcStageMask = mPipelines[i]->mSrcStageMask;
		subpassDependencys[i + 1].srcAccessMask = mPipelines[i]->mSrcAccessMask;
	}


	VkRenderPassCreateInfo renderPassCreateInfo{};
	renderPassCreateInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
	renderPassCreateInfo.attachmentCount = static_cast<uint32>(attachmentDescriptions.size());
	renderPassCreateInfo.pAttachments = attachmentDescriptions.data();
	renderPassCreateInfo.subpassCount = static_cast<uint32>(subpassDescriptions.size());
	renderPassCreateInfo.pSubpasses = subpassDescriptions.data();
	renderPassCreateInfo.dependencyCount = static_cast<uint32>(subpassDependencys.size());
	renderPassCreateInfo.pDependencies = subpassDependencys.data();

	VK_CHECK_RESULT(vkCreateRenderPass(GVulkanRHI->GetLogicalDevice(), &renderPassCreateInfo, nullptr, &mRenderPass));


	for (uint32 i = 0; i < mPipelines.size(); ++i)
	{
		mPipelines[i]->SetRenderPass(mRenderPass, i);
		mPipelines[i]->Create();
	}

	// Do not keep the renderTarget for GC
	mPipelines.clear();
	mRenderTargets.clear();
}
